[Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8 by 1am9trash · Pull Request #6042 · ROCm/aiter

1am9trash · 2026-10-01T07:09:03Z

Motivation

#5833 runs the sparse MLA decode loop peeled on 64-bit gathers (KV pool > 2 GiB) when split-K leaves each program only a few tiles. That is the right choice on triton 3.8, but on triton 3.7 the peeled kernel spills ~110 VGPRs to scratch.

Technical Details

Add a module-level _TRITON_GE_38 check, using packaging.version like other version checks in aiter.
On 64-bit gathers, use the peeled loop only when triton >= 3.8 and split-K is short. Otherwise always unpeel. The buffer-load path is unchanged.
No behavior change on triton >= 3.8.

Test Plan

Triton 3.7:

bs	tokens	r2 main (µs)	r2 This PR (µs)	r2 speedup	r1 main (µs)	r1 This PR (µs)	r1 speedup
1	5	14.9	14.0	1.07x	18.1	14.6	1.24x
2	10	17.3	16.1	1.07x	20.2	15.5	1.31x
4	20	21.6	18.0	1.20x	28.3	18.2	1.55x
8	40	33.5	26.6	1.26x	39.9	26.6	1.50x
12	60	34.1	28.1	1.21x	41.6	28.0	1.49x
16	80	37.7	33.3	1.13x	45.8	32.9	1.39x
24	120	43.7	39.7	1.10x	52.2	39.7	1.31x

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

github-actions · 2026-10-01T07:09:37Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
`ci:triton-355`	Run the full Triton test suite on MI35X, not only the tests the change affects
`multigpu`	Aiter multi-GPU tests on the 8-GPU runner
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6042 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

Copilot

Copilot review overview

🟡 Changes recommended

The version check incorrectly treats Triton 3.8 prereleases as the fixed 3.8 release.

Review effort: Balanced
Findings: 1

Open (1)

Preserve prerelease markers when checking the fixed version · New

What changed in this PR

Updates sparse MLA dispatch to avoid register spilling on Triton versions below 3.8.

Changes:

Detects the installed Triton version.
Unpeels short split-K loops for 64-bit gathers on older Triton releases.

File	Description
`aiter/ops/triton/attention/pa_decode_sparse.py`	Adds version-gated loop selection.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

1am9trash · 2026-10-02T02:40:51Z

The triton tests all passed except shard 6. Seem to be a upstream issue. I also see the same error in other PR.
Same error: https://github.com/ROCm/aiter/actions/runs/36847007242/job/110320157507?pr=6048

vgokhale

Why can't you update your triton whl instead? 3.7 is now quite old.

1am9trash · 2026-10-06T07:05:41Z

Hi @vgokhale,
Agreed, but the sglang rocm image is still on triton 3.7, and there are lots of triton kernels in it, so a version bump will take some time.
Would it be ok to take this as a temporary guard so 3.7 works well?

vgokhale · 2026-10-12T20:34:33Z

@1am9trash I'm approving this as an exception. I don't want to make a habit of having to support old Triton. Please update to 3.8.

1am9trash · 2026-10-13T06:21:17Z

Hi @vgokhale,
Thanks a lot.
I just fixed the conflict. No PR modification change.

azaidy

LGTM! as exception

ROCm/aiter#6042, sgl-project#5967 and sgl-project#6007 are all ancestors of the AITER pin from sgl-project#44379, so the sgl-project#6042 cherry-pick conflicts in pa_decode_sparse.py and fails the AITER step on gfx942/gfx950. The sgl-project#5967 CSV diff and the group32 other=127 sed (whose file no longer exists) are stale for the same reason.

ROCm/aiter#6042, #5967 and #6007 are all ancestors of the AITER pin from #44379, so the #6042 cherry-pick conflicts in pa_decode_sparse.py and fails the AITER step on gfx942/gfx950. The #5967 CSV diff and the group32 other=127 sed (whose file no longer exists) are stale for the same reason.

…e commit The previous commit dropped the ROCm/aiter#6042 cherry-pick on the wrong premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the sgl-project#6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of sgl-project#6159 and applies cleanly to the pin, so cherry-pick that instead. It is on main, so the PR-ref fetch is no longer needed. sgl-project#5967 and sgl-project#6007 are in the pin, so their patches stay removed.

The previous commit dropped the ROCm/aiter#6042 cherry-pick on the wrong premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the #6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of #6159 and applies cleanly to the pin, so cherry-pick that instead. It is on main, so the PR-ref fetch is no longer needed. #5967 and #6007 are in the pin, so their patches stay removed.

… commit sgl-project#44534 dropped the sgl-project#6042 cherry-pick on the premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the sgl-project#6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of sgl-project#6159 and applies cleanly to the pin, and it is on main, so the PR-ref fetch is no longer needed.

… commit #44534 dropped the #6042 cherry-pick on the premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the #6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of #6159 and applies cleanly to the pin, and it is on main, so the PR-ref fetch is no longer needed.

Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8

042a289

1am9trash requested a review from a team October 1, 2026 07:09

github-actions Bot changed the title ~~Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8~~ [Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8 Oct 1, 2026

github-actions Bot added the Triton/Gluon label Oct 1, 2026

Boss2002n requested a balanced review from Copilot October 1, 2026 09:33

Copilot started reviewing on behalf of Boss2002n October 1, 2026 09:34

View session

Copilot AI reviewed Oct 1, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/pa_decode_sparse.py Resolved

Fix version comparison for triton version check

1cbfb85

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

zufayu requested review from a team and vgokhale October 2, 2026 00:33

zufayu assigned vgokhale Oct 2, 2026

michaelzhang-ai mentioned this pull request Oct 2, 2026

[AMD] Skip the AITER #6042/#5967 patches on gfx1250 sgl-project/sglang#42219

Merged

vgokhale requested changes Oct 5, 2026

View reviewed changes

vgokhale previously approved these changes Oct 12, 2026

View reviewed changes

Merge branch 'main' into v41-decode-triton-3.7

8672515

1am9trash dismissed vgokhale’s stale review via 8672515 October 13, 2026 06:19

kkHuang-amd mentioned this pull request Oct 13, 2026

[AMD][DSV4][8/8] gfx950 ROCm images: pin aiter to 57e2c6b42e sgl-project/sglang#44035

Closed

5 tasks

azaidy approved these changes Oct 13, 2026

View reviewed changes

azaidy merged commit 57e2c6b into ROCm:main Oct 13, 2026

kkHuang-amd mentioned this pull request Oct 15, 2026

[AITER-Upgrade] PR readiness sgl-project/sglang#21302

Open

bingxche mentioned this pull request Oct 15, 2026

[CI Monitor] Daily Report - 2026-10-15 bingxche/sglang-ci-bot#227

Open

michaelzhang-ai mentioned this pull request Oct 15, 2026

[Fix] [AMD] docker: Repair the AITER build step after the 2bb549102 pin bump sgl-project/sglang#44534

Merged

michaelzhang-ai mentioned this pull request Oct 15, 2026

[Fix] [AMD] docker: Restore the AITER sparse MLA spill patch on the 2bb549102 pin sgl-project/sglang#44539

Closed

Navigation Menu

Repository navigation

[Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8 - #6042

[Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8#6042
azaidy merged 3 commits into
ROCm:mainfrom
1am9trash:v41-decode-triton-3.7

1am9trash commented Oct 1, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

1am9trash commented Oct 2, 2026 •

edited

Loading

Uh oh!

vgokhale left a comment

Uh oh!

1am9trash commented Oct 6, 2026

Uh oh!

vgokhale commented Oct 12, 2026

Uh oh!

1am9trash commented Oct 13, 2026 •

edited

Loading

Uh oh!

azaidy left a comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

1am9trash commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Oct 1, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

1am9trash commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vgokhale left a comment

Choose a reason for hiding this comment

Uh oh!

1am9trash commented Oct 6, 2026

Uh oh!

vgokhale commented Oct 12, 2026

Uh oh!

1am9trash commented Oct 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

azaidy left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

1am9trash commented Oct 1, 2026 •

edited

Loading

1am9trash commented Oct 2, 2026 •

edited

Loading

1am9trash commented Oct 13, 2026 •

edited

Loading